import re
import openpyxl
import textdistance
import os
import tqdm
import pickle
import numpy as np
from elasticsearch import Elasticsearch
import sys

import argparse
from zhuziyulei import timeit

p_punct = re.compile("[，。：“”、；？！‘’]+")
p_bracket = re.compile("\{(.*?)\}")


def build_es_index():
    print("building es index")
    es = Elasticsearch(
        "http://localhost:9200",
        timeout=3600
    )
    for fname in ['二程外書_(四庫全書本)_全覽.txt', '二程文集_(四庫全書本)_全覽.txt',
                  '二程遺書_(四庫全書本)_全覽.txt', '程氏經説_(四庫全書本)_全覽.txt']:
        if not es.indices.exists(index=fname[2:4]):
            es.indices.create(index=fname[2:4])
            read_siku(os.path.join("ercheng_data", fname), fname, es)
    return es


@timeit
def read_siku(path, fname, es):
    print("reading siku", fname)
    lines = open(path).readlines()
    for line_num, line in enumerate(lines):
        if len(line.strip()) > 0:
            es.index(index=fname[2:4], id=line_num, body={"key": line.strip()})


def read_ercheng(name, es):
    print("processing yinyong")
    # 创建工作簿
    wb = openpyxl.Workbook()
    # 选择活动工作表
    ws = wb.active
    # 为工作表命名
    ws.title = "Sheet1"
    # xlsx格式文件需要使用openpyxl包读取
    df = openpyxl.load_workbook(name)
    table = df.worksheets[0]  # 取出第一个工作表
    paragraphs = {}
    for line_num, row in enumerate(table.rows):
        # 对每一行数据迭代
        print(line_num)
        if line_num == 0:
            # 第一行是数据名
            col_names = [col.value for col in row]
            col_names.insert(4, "匹配")
            ws.append(col_names)
            continue
        record = [col.value for col in row]

        # 匹配时只需要看大字这一列即可
        paragraph = record[3]
        sentence = record[4]
        book = record[0][:2]
        if paragraph in paragraphs:
            siku_paragraph, best_key = paragraphs[paragraph]
            siku_sentence = lcs_path(sentence, siku_paragraph, do_paragraph=False)
        else:
            cleaned_paragraph = re.sub(p_punct, '', paragraph)
            assert es.indices.exists(index=book)
            result = es.search(index=book, query={"match": {"key": cleaned_paragraph}}, size=3)
            max_score = -1
            best_key = None
            for hit in result["hits"]["hits"]:
                cand_text = hit["_source"]["key"]
                sim_score = textdistance.lcsseq.similarity(cand_text, cleaned_paragraph)
                if max_score < sim_score:
                    # 找相似度最高的
                    max_score = sim_score
                    best_key = cand_text
            siku_paragraph = punctuation_match(paragraph, best_key)
            siku_sentence = lcs_path(sentence, best_key, do_paragraph=False)
            paragraphs[paragraph] = siku_paragraph, best_key
        ws.append(record[:3] + [siku_paragraph, best_key, siku_sentence] + record[5:])
    wb.save("guolu.xlsx")


def process_ercheng(name):
    print("processing yinyong")
    # 创建工作簿
    wb = openpyxl.Workbook()
    # xlsx格式文件需要使用openpyxl包读取
    df = openpyxl.load_workbook(name)
    all_sheets = df.worksheets
    for table in all_sheets:
        # 选择活动工作表
        ws = wb.create_sheet(table.title)
        print(table.title)
        for line_num, row in enumerate(table.rows):
            # 对每一行数据迭代
            print(line_num)
            if line_num == 0:
                # 第一行是数据名
                col_names = [col.value for col in row]
                col_names.insert(2, "过录")
                ws.append(col_names)
                continue
            record = [col.value for col in row]

            # 匹配时只需要看大字这一列即可
            siku_text = record[0]
            if siku_text is None or len(record) < 2:
                ws.append(record[:2] + [siku_text] + record[2:])
                continue
            if table.title == "《经说》":
                zhonghua_text = record[2]
            else:
                zhonghua_text = record[1]
            bracket_contents = re.findall(p_bracket, siku_text)
            bracket_free_siku = re.sub(p_bracket, '{}', siku_text)
            cleaned_siku = re.sub(p_punct, '', bracket_free_siku)
            if zhonghua_text is None or len(zhonghua_text) == 0:
                siku_result = siku_text
            else:
                cleaned_zhonghua = re.sub(p_bracket, '{}', zhonghua_text)
                siku_result = lcs_path(cleaned_zhonghua, cleaned_siku, do_paragraph=True)
                if len(bracket_contents) > 0:
                    # 处理括号内的内容
                    i = 0
                    result = []
                    start = 0
                    for match in re.finditer(r"\{\}", siku_result):
                        result.append(siku_result[start:match.start()])
                        result.append('{' + bracket_contents[i] + '}')
                        start = match.end()
                        i += 1
                    assert i == len(bracket_contents), (
                        i, len(bracket_contents), siku_text, cleaned_siku, zhonghua_text, siku_result)
                    result.append(siku_result[start:])
                    siku_result = ''.join(result)
            ws.append(record[:2] + [siku_result] + record[2:])
    wb.save("guolu.xlsx")


def punctuation_match(gold, unseg):
    i, j = 0, 0
    result = ""
    while i < len(gold) and j < len(unseg):
        if gold[i] == unseg[j]:
            result += unseg[j]
            i += 1
            j += 1
        elif re.match(p_punct, gold[i]):
            result += gold[i]
            i += 1
        else:
            break
    if i >= len(gold) and j >= len(unseg):
        return result
    return lcs_path(gold, unseg, do_paragraph=True)


def lcs_path(gold, unseg, do_paragraph: bool):
    """
    根据lcs获得过录文本，
    :param gold:
    :param unseg:
    :param do_paragraph: 主要看是unseg整段还是单句，区别在于是否要把unseg前面的部分加进来
    :return:
    """
    max_sim = 0
    dp = np.zeros((len(gold) + 1, len(unseg) + 1))
    path = [[(-1, -1)] * (len(unseg) + 1) for _ in range(len(gold) + 1)]
    for i in range(1, len(gold) + 1):
        for j in range(1, len(unseg) + 1):
            if gold[i - 1] == unseg[j - 1]:
                dp[i][j] = dp[i - 1][j - 1] + 1
                path[i][j] = (i - 1, j - 1)
            else:
                if dp[i - 1][j] >= dp[i][j - 1]:
                    dp[i][j] = dp[i - 1][j]
                    path[i][j] = (i - 1, j)
                else:
                    dp[i][j] = dp[i][j - 1]
                    path[i][j] = (i, j - 1)
            max_sim = max(max_sim, dp[i][j])
    i = len(gold)
    j = len(unseg)
    result = []
    assert path[i][j] != (-1, -1)
    while i > 0 and j > 0:
        x, y = path[i][j]
        if x == i - 1 and y == j - 1:
            # s1[i-1] == s2[j-1]
            result.append(unseg[j - 1])
        elif x == i - 1 and y == j:
            # gold[i-1]被跳过，如果是标点需要保留，否则跳过
            if gold[i - 1] in "，。；：！？、“”‘’":
                result.append(gold[i - 1])
        elif x == i and y == j - 1:
            # unseg[j-1]被跳过，需要保留
            result.append(unseg[j - 1])
        i, j = x, y
    result = result[::-1]
    result = "".join(result)
    if j > 0 and do_paragraph:
        result = unseg[:j] + result
    if not do_paragraph and i > 0:
        result = gold[:i] + result
    return result


if __name__ == '__main__':
    process_ercheng("ercheng_new.xlsx")
